package cn.xiaolang.function.data_structure;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 限流算法-漏桶算法
 * <p>
 * 漏桶算法能够强行限制数据的传输速率。比如单机设置不能超过 100qps，当超过时就拒绝服务
 * <p>
 * 缺点：
 * 漏出速率是固定，对突发流量来说不是很好
 *
 * @author fuhw/Dean
 * @date 2020-04-10
 */
public final class FlowLimitByLeakyBucket {

    // 桶的容量
    private int capacity;
    // 木桶剩余的水滴的量(初始化的时候的空的桶)
    private AtomicInteger water = new AtomicInteger(0);
    // 第一次请求之后,木桶在这个时间点开始漏水
    private long leakTimestamp = 0;
    // 水滴的流出的速率 每秒流出频率
    private int leakRatePerSecond;


    private FlowLimitByLeakyBucket(int capacity, int leakRatePerSecond) {
        this.capacity = capacity;
        this.leakRatePerSecond = leakRatePerSecond;
    }

    public static FlowLimitByLeakyBucket build(int capacity, int leakRatePerSecond) {
        return new FlowLimitByLeakyBucket(capacity, leakRatePerSecond);
    }

    public static FlowLimitByLeakyBucket build(int capacity) {
        return new FlowLimitByLeakyBucket(capacity, capacity);
    }

    private boolean acquire() {
        // 如果是空桶，就当前时间作为桶开是漏出的时间
        if (water.get() == 0) {
            leakTimestamp = System.currentTimeMillis();
            water.getAndIncrement();
            return true;
        }

        // 计算剩余水量
        int waterLeft = water.get() - (int) ((System.currentTimeMillis() - leakTimestamp) * leakRatePerSecond / 1000);
        water.set(Math.max(0, waterLeft));
        // 尝试加水,并且水还未满
        if (water.get() < capacity) {
            // 更新leakTimeStamp
            leakTimestamp = System.currentTimeMillis();
            water.getAndIncrement();
            return true;
        }
        return false;
    }

    public static void main(String[] args){
        FlowLimitByLeakyBucket limit = FlowLimitByLeakyBucket.build(100);
        AtomicLong successCount = new AtomicLong();
        AtomicLong totalCount = new AtomicLong();
        long start = System.currentTimeMillis();
        while (true) {
            if (limit.acquire()) {
                successCount.getAndIncrement();
            }
            totalCount.getAndIncrement();

            long consume = (System.currentTimeMillis() - start) / 1000L;
            long qps = consume == 0 ? totalCount.get() : totalCount.get() / consume;
            System.out.println("total:" + totalCount.get() +
                    ",success:" + successCount.get() +
                    ",consume:" + consume + "(s)" +
                    ",qps:" + qps + "/s");
        }
    }
}
